# -*- coding: utf-8 -*-
# @Time    : 2021/7/7 14:46
# @Author  : LuoTianHang


# ####################metric.py 说明##########################
# this script is used to metric the performance of the model
import numpy as np

eps = 1e-5


class PCKh:
    def __init__(self):
        self.threshold_distance = 0.5
        self.joints = 16

        self.total = [0 + eps] * self.joints
        self.correct = [0] * self.joints

    def update(self, predictions, targets, vis, scale):
        predictions = PCKh.get_preds(predictions.cpu().numpy())
        targets = PCKh.get_preds(targets.cpu().numpy())
        vis = vis.cpu().numpy()
        scale = scale.cpu().numpy()
        self.cal(predictions, targets, vis, scale)

    def get(self):

        result = {
            "ankles": round((self.correct[0] + self.correct[5]) / (self.total[0] + self.total[5]), 3),
            "knee": round((self.correct[1] + self.correct[4]) / (self.total[1] + self.total[4]), 3),
            "hip": round((self.correct[2] + self.correct[3]) / (self.total[2] + self.total[3]), 3),
            "pelvis": round(self.correct[6] / self.total[6], 3),
            "thorax": round(self.correct[7] / self.total[7], 3),
            "upper neck": round(self.correct[8] / self.total[8], 3),
            "head top": round(self.correct[9] / self.total[9], 3),
            "wrist": round((self.correct[10] + self.correct[15]) / (self.total[10] + self.total[15]), 3),
            "elbow": round((self.correct[11] + self.correct[14]) / (self.total[11] + self.total[14]), 3),
            "shoulder": round((self.correct[12] + self.correct[13]) / (self.total[12] + self.total[13]), 3),
            "total": round(sum(self.correct) / sum(self.total), 3)
        }

        return result

    def reset(self):
        self.total = [0 + eps] * self.joints
        self.correct = [0] * self.joints

    def cal(self, predictions, targets, vis, scales):
        for i in range(predictions.shape[0]):
            scale = scales[i]

            for index, (p, d, v) in enumerate(zip(predictions[i], targets[i], vis[i])):
                if v == 0:
                    pass
                else:
                    self.total[index] += 1
                    dis = np.linalg.norm(p - d, 2)  # pred与标注的欧式距离
                    if dis < self.threshold_distance * scale:
                        self.correct[index] += 1

    @staticmethod
    def get_preds(maps):
        res = maps.shape[2]
        maps = maps.reshape(maps.shape[0], maps.shape[1], maps.shape[2] * maps.shape[3])
        idx = np.argmax(maps, axis=2)
        preds = np.zeros((maps.shape[0], maps.shape[1], 2))
        preds[:, :, 0], preds[:, :, 1] = idx % res, idx // res
        return preds
